import math
from copy import deepcopy

import numpy as np
from scipy.stats import beta

from PBLTL.language import P
from PBLTL.parser import parse
from model.nptem import NPTEM

class ModelChecker:
    def __init__(self, nptem: NPTEM, pbltl: P):
        self.nptem = nptem
        self.pbltl = pbltl


    def confidenc_smc(self, time_bound:float, theta:float, coverage:float ) :
        """
        The sampling number is calculating using Wald Confidence interval.
        """

        pass

    def apmc_smc(self, time_bound:float, theta:float, coverage:float ) :
        """
        The sampling number is calculating using Chernoff-Hoeffding bound
        """
        traces = []
        counter_traces = []
        n = ch_bound(coverage, theta)
        x = 0
        for i in range(n):
            trace = self.nptem.run(time_bound, mode="stochastic", log_mode=None)
            traces.append(trace)
            phi = self.pbltl.phi
            if phi.eval(trace, 0):
                x += 1
            else:
                counter_traces.append(trace)
        return x / n, n, traces, counter_traces


    def bayesian_smc(self, time_bound:float, theta:float, coverage:float, a:float, b:float, ) :
        # todo 多线程
        n = 0
        x = 0
        traces = []
        counter_traces = []
        while True:
            nptem = deepcopy(self.nptem)
            trace = nptem.run(time_bound, mode="stochastic", log_mode=None)
            traces.append(trace)
            n += 1
            phi = self.pbltl.phi
            if phi.eval(trace, 0):
                x += 1
                print("PBLTL is true in %d times: " % n)
            else:
                print("PBLTL is false in %d times: " % n)
                counter_traces.append(trace)
            p = (x + a) / (n + a + b)
            t0 = p - theta
            t1 = p + theta

            if t1 > 1:
                t0 = 1 - 2 * theta
                t1 = 1
            elif t0 < 0:
                t0 = 0
                t1 = 2 * theta

            alpha_post = a + x
            beta_post = b + (n - x)
            posterior = beta(alpha_post, beta_post)
            gamma = posterior.cdf(t1) - posterior.cdf(t0)
            # gamma = beta.pdf(t1, x + a, n - x + b) - beta.pdf(t0, x + a, n - x + b)
            print(gamma)
            if gamma >= coverage:
                return p, n, traces, counter_traces


def ch_bound(coverage, theta):
    delta = 1 - coverage
    n = np.log(2 / delta) / (2 * theta ** 2)
    return math.ceil(n)


if __name__ == "__main__":
    nptem = NPTEM.mlag()
    s = "P=?{n1.state=up U[10000, inf) n1.state=down}"
    time_bound = 1000 * 365 * 24 * 60 * 60
    pbltl = parse(s)
    theta = 0.05
    coverage = 0.95 # coverage of the model checker

    mc = ModelChecker(nptem, pbltl)
    print(ch_bound(coverage, theta))
    p, n, _, counter_traces = mc.bayesian_smc(time_bound, theta, coverage, 1, 1)
    print(p, n, theta)
    print(counter_traces)


    # p, n, _, counter_traces = mc.apmc_smc(time_bound, theta, coverage)
    # print(p, n, theta)
    # print(counter_traces)
